import time
import sentencepiece


trainer = sentencepiece.SentencePieceTrainer

start = time.perf_counter()
# 参数 https://github.com/google/sentencepiece/blob/master/doc/options.md
trainer.train(
    input="data/processed/peoples_daily.txt",  # 训练语料，形式为一行一段文本
    model_prefix="lulu/tokenizer",             # 模型输出的路径
    shuffle_input_sentence=False,              # 是否打乱句子
    max_sentence_length=16384,                 # 句子最大长度
    pad_id=3,                                  # PAD id
    model_type="bpe",                          # 训练时模型的类别：unigram (default), bpe, char, word
    vocab_size=60000,                          # 训练后词表的大小，数量越大训练越慢
    split_digits=True,                         # 将所有数字 (0-9) 分成单独的部分
    split_by_unicode_script=True,              # 使用 Unicode 脚本分割句子片段
    byte_fallback=True,                        # 将未知片段分解为 UTF-8 字节片段
    allow_whitespace_only_pieces=True,         # 允许仅包含（连续）空白标记的片段
    remove_extra_whitespaces=False,            # 删除前导、尾随和重复的内部空格
    normalization_rule_name="nfkc",            # 规范化规则名称。nfkc 或 identity
    num_threads=128                            # 训练线程数
)
end = time.perf_counter()
print("elapsed time:", end - start)
